
from qiskit import QuantumCircuit, ClassicalRegister, transpile
from qiskit_aer import AerSimulator
import numpy as np
from typing import Callable, List, Union, Optional,Any
from collections.abc import Iterator
import random
import itertools

from algorithm.grover_operator import grover_operator

# 定义幅度放大问题类
class AmplificationProblem:
    def __init__(
        self,
        oracle: QuantumCircuit,
        state_preparation: QuantumCircuit | None = None,
        objective_qubits: Union[int, List[int], None] = None,
        is_good_state: Callable[[str], bool] | None = None,
        post_processing: Callable[[str], Any] | None = None
    ):
        """
        定义幅度放大问题。

        参数:
            oracle: 相位 Oracle 电路
            state_preparation: 状态准备电路 A，若为 None 则使用全 H 门
            objective_qubits: 测量量子比特，若为 None 则测量所有量子比特
            is_good_state: 判断好态的函数
            post_processing: 后处理函数，默认为恒等映射
        """
        self._oracle = oracle
        self._state_preparation = state_preparation
        self._objective_qubits = objective_qubits
        self._is_good_state = is_good_state
        self._post_processing = post_processing

    @property
    def oracle(self) -> QuantumCircuit:
        """返回 Oracle 电路"""
        return self._oracle

    @property
    def state_preparation(self) -> QuantumCircuit:
        """返回状态准备电路 A，若未指定则使用全 H 门"""
        if self._state_preparation is None:
            circuit = QuantumCircuit(self.oracle.num_qubits)
            circuit.h(range(self.oracle.num_qubits))
            return circuit
        return self._state_preparation

    @property
    def objective_qubits(self) -> List[int]:
        """返回测量量子比特的索引列表"""
        if self._objective_qubits is None:
            return list(range(self.oracle.num_qubits))
        if isinstance(self._objective_qubits, int):
            return [self._objective_qubits]
        return self._objective_qubits

    @property
    def is_good_state(self) -> Callable[[str], bool]:
        """返回判断好态的函数，若未提供则抛出错误"""
        if self._is_good_state is None:
            raise ValueError("必须提供 is_good_state 函数")
        return self._is_good_state

    @property
    def post_processing(self) -> Callable[[str], Any]:
        """返回后处理函数，默认为恒等映射"""
        if self._post_processing is None:
            return lambda x: x
        return self._post_processing

# 定义结果类
class GroverResult:
    def __init__(self):
        """初始化结果对象，存储算法运行结果"""
        self.iterations = None
        self.top_measurement = None
        self.assignment = None
        self.oracle_evaluation = None
        self.circuit_results = None
        self.max_probability = None

# Grover 算法类
class Grover:
    def __init__(
        self,
        iterations: Union[int, List[int], Iterator[int], None] = None,
        growth_rate: float | None = None,
        sample_from_iterations: bool = False,
        shots: int = 1024
    ):
        """
        初始化 Grover 算法。

        参数:
            iterations: Grover 算子迭代次数（整数、列表、迭代器或 None）
            growth_rate: 若指定，迭代次数按增长率递增
            sample_from_iterations: 是否从迭代次数中随机采样
            shots: 模拟器运行的测量次数
        """
        if growth_rate is not None and iterations is not None:
            raise ValueError("不能同时指定 iterations 和 growth_rate")
        if growth_rate is not None:
            if growth_rate <= 1:
                raise ValueError("growth_rate 必须大于 1")
            self._iterations = (int(growth_rate**x) for x in itertools.count(1))
        elif isinstance(iterations, int):
            self._iterations = [iterations]
        else:
            self._iterations = iterations if iterations is not None else []
        self._sample_from_iterations = sample_from_iterations
        self._shots = shots
        self._simulator = AerSimulator()

    @staticmethod
    def optimal_num_iterations(num_solutions: int, num_qubits: int) -> int:
        """
        计算已知解数量时的最优迭代次数。

        参数:
            num_solutions: 解的数量
            num_qubits: 量子比特数量

        返回:
            最优迭代次数
        """
        amplitude = np.sqrt(num_solutions / 2**num_qubits)
        return round(np.arccos(amplitude) / (2 * np.arcsin(amplitude)))

    def construct_circuit(
        self,
        problem: AmplificationProblem,
        power: int,
        measurement: bool = False
    ) -> QuantumCircuit:
        """
        构造 Grover 算法电路。

        参数:
            problem: 幅度放大问题
            power: Grover 算子重复次数
            measurement: 是否添加测量

        返回:
            QuantumCircuit: 构造好的电路
        """
        qc = QuantumCircuit(problem.oracle.num_qubits, name="Grover circuit")
        qc.compose(problem.state_preparation, inplace=True)
        if power > 0:
            grover_op = grover_operator(problem.oracle, problem.state_preparation)
            qc.compose(grover_op.power(power), inplace=True)
        if measurement:
            measurement_cr = ClassicalRegister(len(problem.objective_qubits))
            qc.add_register(measurement_cr)
            qc.measure(problem.objective_qubits, measurement_cr)
        return qc

    def amplify(self, problem: AmplificationProblem) -> GroverResult:
        """
        运行幅度放大算法。

        参数:
            problem: 幅度放大问题

        返回:
            GroverResult: 算法结果
        """
        result = GroverResult()
        iterations = []
        all_circuit_results = []
        max_probability = 0
        top_measurement = "0" * len(problem.objective_qubits)
        oracle_evaluation = False

        max_iterations = max(10, 2**problem.oracle.num_qubits)
        max_power = np.ceil(2 ** (len(problem.objective_qubits) / 2))

        if isinstance(self._iterations, list):
            iterator = iter(self._iterations)
        else:
            iterator = self._iterations

        for _ in range(max_iterations):
            try:
                power = next(iterator)
            except StopIteration:
                break
            if power > max_power:
                break
            iterations.append(power)
            if self._sample_from_iterations:
                power = random.randint(0, power)
            qc = self.construct_circuit(problem, power, measurement=True)
            qc_transpiled = transpile(qc, self._simulator)
            job = self._simulator.run(qc_transpiled, shots=self._shots)
            counts = job.result().get_counts()
            circuit_results = {label: count / self._shots for label, count in counts.items()}
            all_circuit_results.append(circuit_results)
            top_measurement, max_probability = max(circuit_results.items(), key=lambda x: x[1])
            if problem.is_good_state is not None:
                oracle_evaluation = problem.is_good_state(top_measurement)
                if oracle_evaluation:
                    break

        result.iterations = iterations
        result.top_measurement = top_measurement
        result.assignment = problem.post_processing(top_measurement)
        result.oracle_evaluation = oracle_evaluation
        result.circuit_results = all_circuit_results
        result.max_probability = max_probability
        return result

def make_oracle_for_11(num_qubits=2):
    """构造一个标记 |11> 的 Oracle"""
    qc = QuantumCircuit(num_qubits, name="Oracle")
    # 对 |11> 添加相位翻转
    qc.cz(0, 1)
    return qc


def test_grover_operator_basic():
    """验证 grover_operator 结构正确性"""
    oracle = make_oracle_for_11()
    state_prep = QuantumCircuit(2)
    state_prep.h([0, 1])

    grover_op = grover_operator(oracle, state_prep)
    assert isinstance(grover_op, QuantumCircuit)
    # 电路至少包含 oracle 和一轮反射
    assert grover_op.num_qubits == 2
    assert grover_op.size() > oracle.size()


def test_amplification_problem_properties():
    """验证 AmplificationProblem 属性是否正常工作"""
    oracle = make_oracle_for_11()
    problem = AmplificationProblem(
        oracle=oracle,
        is_good_state=lambda x: x == "11"
    )

    # 状态准备默认全 H
    sp = problem.state_preparation
    assert isinstance(sp, QuantumCircuit)
    assert sp.size() > 0

    # objective_qubits 默认全体
    assert problem.objective_qubits == [0, 1]

    # good_state 判断函数能用
    assert problem.is_good_state("11")
    assert not problem.is_good_state("00")

    # post_processing 默认恒等
    assert problem.post_processing("11") == "11"


def test_grover_end_to_end():
    """端到端测试 Grover 算法"""
    oracle = make_oracle_for_11()
    problem = AmplificationProblem(
        oracle=oracle,
        is_good_state=lambda x: x == "11"
    )
    grover = Grover(iterations=[1, 2, 3], shots=1024)

    result = grover.amplify(problem)

    # 检查结果对象
    assert result.iterations is not None
    assert result.top_measurement is not None
    assert result.circuit_results is not None
    assert result.max_probability is not None

    # |11> 应该有较高概率（通常 >0.6）
    assert result.top_measurement == "11"
    assert result.max_probability > 0.6


def test_optimal_num_iterations():
    """验证最优迭代次数计算"""
    num_qubits = 2
    num_solutions = 1
    iters = Grover.optimal_num_iterations(num_solutions, num_qubits)
    # 对 2 qubit，目标态数量=1，最优迭代次数应为 1
    assert iters == 1

def make_oracle_for_111(num_qubits=3):
    """构造一个标记 |111> 的 Oracle"""
    qc = QuantumCircuit(num_qubits, name="Oracle")
    # 多控 Z，相当于对 |111> 添加相位翻转
    qc.h(2)           # 将 Z 作用转换为受控 X
    qc.ccx(0, 1, 2)   # Toffoli 门
    qc.h(2)
    return qc


def test_grover_three_qubits():
    """测试 3-qubit Grover 算法，目标态 |111>"""
    oracle = make_oracle_for_111(3)
    problem = AmplificationProblem(
        oracle=oracle,
        is_good_state=lambda x: x == "111"
    )

    grover = Grover(iterations=[1, 2, 3], shots=2048)
    result = grover.amplify(problem)

    assert result.top_measurement == "111"
    # 对 3 qubit，一般迭代 1-2 次即可把成功率拉高到 >0.5
    assert result.max_probability > 0.5
